import torch
from torch import nn
from tabulate import tabulate
from transformers.utils.fx import symbolic_trace
from torch.fx.passes.shape_prop import ShapeProp

__doc__ = """
use transformers.utils.fx(trace transformers models) to get static graph
Author: jirongting
datetime: 2024.04.15
version: 2 2024.04.15 second commit
"""

"""
check if list(fx.Node.meta) has tensorMeta
Parameters:
    nodeMeta - fx.Node.meta which contains the shape of ops
"""

def is_tensor(nodeMeta):
    if isinstance(nodeMeta, dict):
        for key,value in nodeMeta.items():
            if isinstance(value,(torch._C._TensorMeta,torch.fx.passes.shape_prop.TensorMetadata)):# if the (key,value) has attr tensor
                return True
            elif isinstance(value,dict):
                if is_tensor(value):
                    return True
    return False

"""
add inputShape attr to node.meta

Parameters:
 traced - the model traced by fx
"""

def addInputShapeToNode(traced):
    for node in traced.graph.nodes:
        input_shape=[] # init the input_shape
        for i in node._input_nodes: # add input shape of node
            if is_tensor(i.meta): # if the input node has output, add the shape into input_shape else [] 
                input_shape.append(i.meta['tensor_meta'].shape)
            else:
                input_shape.append([])
        node.meta['InputShape']=input_shape
        
""" 
use transformers.utils.fx  and ShapeProp(torch.Interpreter) to get the static graph of model  

Parameters: 
 model - nn.Module 
 input_size - the input size which help to get shape of each node
Returns: 
 return a static graph of model 
"""
def getTorchGraph(model:nn.Module,input_size:torch.Tensor,pre = True):
    if pre:
        traced = symbolic_trace(model) # use transformers.utils.fx to trace the model
    else:
        dummy_input_embeds = {
            'inputs_embeds': input_size
        }
        traced = symbolic_trace(model,input_names=list(dummy_input_embeds.keys()))
    ShapeProp(traced).propagate(input_size) # use ShapeProp to get shape of each node
    addInputShapeToNode(traced)
    return traced
